from sklearn.datasets import make_moons
from matplotlib import pyplot
from matplotlib.colors import ListedColormap
import numpy as np
import os
from origin import source
def main():
    np.random.seed(47)
    # np.random.seed(47)
    import shutil
    shutil.rmtree('origin_scalar')
    os.mkdir('origin_scalar')
    # Xs,Ys,Xt,Yt = make_trans_moons(theta=30,number=200)
    #200 40 93.5#300 30 96 #200 20 99.5#320 20
    Xs, Ys, Xt, Yt = make_trans_moons(theta=30, number=200)

    origin = source(Xs=Xs,Ys=Ys,Xt=Xt,Yt=Yt,save=True)
    # origin.fit()
    pyplot.figure()
    pyplot.xticks([])
    pyplot.yticks([])
    # pyplot.axis('off')
    pyplot.title("origin classification")\

    draw_trans_data(Xs,Ys,Xt,origin.predict)
    pyplot.savefig('origin.pdf')
    pyplot.show()
    pyplot.close()


def make_trans_moons(theta=30,number=100,noise=.05):
    from math import cos,sin,pi
    Xs,Ys = make_moons(number,noise=noise,random_state=1)
    Xt,Yt = make_moons(number,noise=noise,random_state=2)

    trans = -np.mean(Xs,axis=0)
    Xs = 2*(Xs + trans)
    Xt = 2*(Xt + trans)
    #旋转
    theta = -theta*pi/180
    rotation = np.array([[cos(theta),sin(theta)],[-sin(theta),cos(theta)]])
    Xt = np.dot(Xt,rotation.T)
    return Xs,Ys,Xt,Yt


def draw_trans_data(Xs,Ys,Xt,predict=None):
    # cm_bright = ListedColormap(['#FF0000','#33FFFF','#00FF00'])
    cm_bright = ListedColormap(['#FFFC33','#33FFF3'])
    x_min,x_max = 1.1*Xs[:,0].min(),1.1*Xs[:,0].max()
    y_min,y_max = 1.5*Xs[:,1].min(),1.5*Xs[:,1].max()

    pyplot.xlim((x_min,x_max))
    pyplot.ylim((y_min,y_max))

    pyplot.tick_params(direction='in',labelleft=False)

    if predict is not None:
        h=.02
        x,y = np.meshgrid(np.arange(x_min,x_max,h),
                          np.arange(y_min,y_max,h))
        # Z1,Z2 = predict(np.c_[x.ravel(),y.ravel()])
        # Z1 = Z1.reshape(x.shape).numpy()
        # Z2 = Z2.reshape(x.shape).numpy()
        # Z = Z1+Z2
        Z = predict(np.c_[x.ravel(),y.ravel()])
        Z = Z.reshape(x.shape).numpy()
        pyplot.contourf(x,y,Z,cmap=cm_bright,alpha=.4)
        pyplot.contour(x,y,Z,colors='black',linewidths=2)

    if Xs is not None:
        for i in range(len(Ys)):
            if Ys[i] == 1:
                pyplot.annotate(".",Xs[i,:],color="green",size=40*1.5)
            else:
                pyplot.annotate(".",Xs[i,:],color="red",size=40*1.5)

    if Xt is not None:
        pyplot.scatter(Xt[:,0],Xt[:,1],c='b',s=40)

if __name__ == '__main__':
    main()

